
import {Body, Controller, Get, HttpStatus, Post, Query, Res, Headers} from '@nestjs/common';
import {response, Response} from 'express';
import {CommonRedisService} from '../service/common.redis.service';
import {CommonRequestService} from '../service/common.request.service';
import {RecordListService} from '../service/record.list.service';
import {UserListService} from '../service/user.list.service';
import {CommonResponseDto} from '../dto/common.response.dto';
import {ResCode, UserStatus} from '../common/constants';
import AppConfig from 'config/app.config';
import {formatTimeToDateStr, uuid} from "../common/utils";
import {decrypt, getHeadersApiToken, getModelPrice, getPlatformModelObj, setErrorStreamRes} from "../common/common";
import Decimal from "decimal.js";
import {TencentService} from "../service/tencent.service";

const volcesChat = async ({modelObj, content,userId}:{modelObj: {code: string,label: string,model: string,price: any,platformPrice: any}, content: string, userId: string}, self: ChatController, res: Response) => {
    try {
        const modelPrice = {
            inPrice: modelObj.price.in,
            outPrice: modelObj.price.out,
        }
        const useage = {
            outTokens: 0,
            inTokens: 0,
            totalTokens: 0,
        }
        const response = await self.commonRequestService.postStream('https://ark.cn-beijing.volces.com/api/v3/chat/completions', {
            model: modelObj.model,
            messages: [
                {role: 'user', content: content}
            ],
            stream: true,
            stream_options: {include_usage:true}
        }, {
            headers: {
                Authorization: `Bearer ${AppConfig.ApiKeyConfig.volces}`
            },
            responseType: 'stream',
        })
        // @ts-ignore
        if (response.data && typeof response.data.on === 'function') {
            // @ts-ignore
            response.data.on('data', (chunk) => {
                const messageArr = chunk.toString().replace('data: ', '').split('data: ')
                messageArr.forEach((msg) => {
                    if (msg .indexOf('choices') > -1) {
                        try {
                            const data = JSON.parse(msg)
                            const choices = data.choices
                            if (data.usage) {
                                useage.inTokens = data.usage.prompt_tokens
                                useage.outTokens = data.usage.completion_tokens
                                useage.totalTokens = data.usage.total_tokens
                            }
                            if (choices.length > 1) {
                                choices.sort((a, b) => {
                                    return a.index-b.index
                                })
                            }
                            let obj: any = {
                                id: data.Id
                            }
                            choices.forEach((choice) => {
                                const delta = choice.delta
                                if (delta.hasOwnProperty('reasoning_content')) {
                                    if (!obj.reasoning) {
                                        obj.reasoning = ''
                                    }
                                    obj.reasoning += delta.reasoning_content
                                } else if (delta.hasOwnProperty('content')) {
                                    if (!obj.answer) {
                                        obj.answer = ''
                                    }
                                    obj.answer += delta.content
                                }
                            })
                            res.write('data: '+JSON.stringify(obj)+'\n\n')
                        }catch(e){
                            console.log(e)
                        }
                    }
                })

            })
            // @ts-ignore
            response.data.on('end', () => {
                console.log('本次回答计量：' + JSON.stringify(useage))
                const date = formatTimeToDateStr(new Date().getTime())
                self.recordListService.createUpdateOne({
                    id: uuid(32),
                    userId: userId,
                    date,
                    inToken: Number(useage.inTokens),
                    outToken: Number(useage.outTokens),
                    totalToken: Number(useage.totalTokens),
                    model: `${modelObj.code}[#]${modelObj.model}`,
                    inPrice: '',
                    outPrice: ''
                }, modelPrice)
                res.end('[DONE]')
            })
        }
    } catch(e){
        console.log(e)
        setErrorStreamRes('服务异常', res)
    }

}
const tencentChat = async ({modelObj, content,userId}:{modelObj: {code: string,label: string,model: string,price: any,platformPrice: any}, content: string, userId: string}, self: ChatController, res: Response) => {
    const modelPrice = {
        inPrice: modelObj.price.in,
        outPrice: modelObj.price.out,
    }
    const useage = {
        outTokens: 0,
        inTokens: 0,
        totalTokens: 0,
    }
    self.tencentService.chatComp(content, modelObj.model, async (msg) => {
        const data = msg.data
        // console.log('message data', data)
        if (data === '[DONE]') {
            console.log('本次回答计量：' + JSON.stringify(useage))
            const date = formatTimeToDateStr(new Date().getTime())
            self.recordListService.createUpdateOne({
                id: uuid(32),
                userId: userId,
                date,
                inToken: Number(useage.inTokens),
                outToken: Number(useage.outTokens),
                totalToken: Number(useage.totalTokens),
                model: `${modelObj.code}[#]${modelObj.model}`,
                inPrice: '',
                outPrice: ''
            }, modelPrice)
            res.end('[DONE]')
        } else {
            try {
                const dataObj = JSON.parse(data)
                useage.outTokens = dataObj.Usage.CompletionTokens
                useage.inTokens = dataObj.Usage.PromptTokens
                useage.totalTokens = dataObj.Usage.TotalTokens
                const choices = JSON.parse(data).Choices
                if (choices.length > 1) {
                    choices.sort((a, b) => {
                        return a.Index-b.Index
                    })
                }
                let obj: any = {
                    id: dataObj.Id
                }
                choices.forEach((choice) => {
                    const delta = choice.Delta
                    if (delta.hasOwnProperty('ReasoningContent')) {
                        if (!obj.reasoning) {
                            obj.reasoning = ''
                        }
                        obj.reasoning += delta.ReasoningContent
                    } else if (delta.hasOwnProperty('Content')) {
                        if (!obj.answer) {
                            obj.answer = ''
                        }
                        obj.answer += delta.Content
                    }
                })
                res.write('data: '+JSON.stringify(obj)+'\n\n')
            } catch(e){
                console.log(e)
            }
        }
    }, () => {
        res.end('')
    })
}
const baiduChat = async ({modelObj, content,userId}:{modelObj: {code: string,label: string,model: string,price: any,platformPrice: any}, content: string, userId: string}, self: ChatController, res: Response) => {
    const modelPrice = {
        inPrice: modelObj.price.in,
        outPrice: modelObj.price.out,
    }
    const useage = {
        outTokens: 0,
        inTokens: 0,
        totalTokens: 0,
    }
    const response = await self.commonRequestService.postStream('https://qianfan.baidubce.com/v2/chat/completions', {
        model: modelObj.model,
        messages: [
            {role: 'user', content: content}
        ],
        stream: true,
        stream_options: {include_usage:true}
    }, {
        headers: {
            'Content-Type': 'application/json',
            appid: AppConfig.ApiKeyConfig.baidu.appId,
            Authorization: `Bearer ${AppConfig.ApiKeyConfig.baidu.apiKey}`
        },
        responseType: 'stream',
    })
    // @ts-ignore
    if (response.data && typeof response.data.on === 'function') {
        // @ts-ignore
        response.data.on('data', (chunk) => {
            const message = chunk.toString().replace('data: ', '')
            if (message .indexOf('choices') > -1) {
                try {
                    const data = JSON.parse(message)
                    const choices = data.choices
                    if (data.usage) {
                        useage.inTokens = data.usage.prompt_tokens
                        useage.outTokens = data.usage.completion_tokens
                        useage.totalTokens = data.usage.total_tokens
                    }
                    if (choices.length > 1) {
                        choices.sort((a, b) => {
                            return a.index-b.index
                        })
                    }
                    let obj: any = {
                        id: data.id
                    }
                    choices.forEach((choice) => {
                        const delta = choice.delta
                        if (delta.hasOwnProperty('reasoning_content')) {
                            if (!obj.reasoning) {
                                obj.reasoning = ''
                            }
                            obj.reasoning += delta.reasoning_content
                        } else if (delta.hasOwnProperty('content')) {
                            if (!obj.answer) {
                                obj.answer = ''
                            }
                            obj.answer += delta.content
                        }
                    })
                    res.write('data: '+JSON.stringify(obj)+'\n\n')
                }catch(e){
                    console.log(e)
                }
            }
        })
        // @ts-ignore
        response.data.on('end', () => {
            console.log('本次回答计量：' + JSON.stringify(useage))
            const date = formatTimeToDateStr(new Date().getTime())
            self.recordListService.createUpdateOne({
                id: uuid(32),
                userId: userId,
                date,
                inToken: Number(useage.inTokens),
                outToken: Number(useage.outTokens),
                totalToken: Number(useage.totalTokens),
                model: `${modelObj.code}[#]${modelObj.model}`,
                inPrice: '',
                outPrice: ''
            }, modelPrice)
            res.end('[DONE]')
        })
    }
}
@Controller('stream')
export class ChatController {
    constructor(private readonly commonRedisService: CommonRedisService, readonly tencentService: TencentService,
                readonly commonRequestService: CommonRequestService,
                readonly recordListService: RecordListService, private readonly userListService: UserListService) {
    }

    @Post('/chat')
    async chat(@Headers() headers: Record<string, string>, @Body() body: {content: string, model: string }, @Res({ passthrough: false }) res: Response) {
        console.log(body)
        res.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
        res.setHeader('Cache-Control', 'no-cache');
        res.setHeader('Connection', 'keep-alive');
        res.status(HttpStatus.OK);
        if (!body.model || !body.content) {
            setErrorStreamRes('参数缺失model或content', res)
        } else {
            const modelConfig = AppConfig.ModelList.find((ml) => ml.code === body.model);
            if (!modelConfig) {
                setErrorStreamRes('参数model错误', res)
            } else {
                const headerToken = getHeadersApiToken(headers)
                const decryptStr = decrypt(headerToken, AppConfig.ApiTokenSecretKey)
                const strArr = decryptStr.split('_')
                if (strArr.length !== 2) {
                    setErrorStreamRes('api-token非法', res)
                } else {
                    const userId: string = strArr[0]
                    const token = await this.commonRedisService.get(AppConfig.ApiHeaderKey + '_' + userId)
                    if (!token || token !== headerToken) {
                        setErrorStreamRes('非法用户', res)
                    } else {
                        const user = await this.userListService.getOneById(userId)
                        if (!user) {
                            setErrorStreamRes('用户不存在', res)
                        } else if (user.status !== UserStatus.NORMAL) {
                            setErrorStreamRes('用户状态异常', res)
                        } else if (new Decimal(user.balance).lessThanOrEqualTo(0.01)) {
                            setErrorStreamRes('余额不足0.01', res)
                        } else {
                            const platformModelObj: {code: string,label: string,model: string,price: any,platformPrice: any} = getPlatformModelObj(modelConfig)
                            console.log('当前对话使用厂商信息：', platformModelObj.label, platformModelObj.model)
                            switch(platformModelObj.code) {
                                case 'tencent':
                                    tencentChat({modelObj: platformModelObj,content: body.content, userId: userId}, this, res)
                                    break;
                                case 'baidu':
                                    baiduChat({modelObj: platformModelObj,content: body.content, userId: userId}, this, res)
                                    break;
                                case 'volces':
                                    volcesChat({modelObj: platformModelObj,content: body.content, userId: userId}, this, res)
                                    break;
                            }
                        }
                    }
                }
            }
        }
    }
}